  module vertical_coord_mod
    use spark_cubed_sphere
    use static_mod
    use params_mod
    use math_mod
    implicit none
    
    interface
      subroutine vertical_distribution_interface(xi,dxideta,detadxi,eta,eta_min,eta_max,z_min,z_max,m_coef)
        import r_kind
        real(r_kind), intent(out) :: xi
        real(r_kind), intent(out) :: dxideta
        real(r_kind), intent(out) :: detadxi
        real(r_kind), intent(in ) :: eta
        real(r_kind), intent(in ) :: eta_min,eta_max
        real(r_kind), intent(in ) :: z_min,z_max
        real(r_kind), intent(in ) :: m_coef
      end subroutine vertical_distribution_interface
      
      subroutine vertical_coordinate_interface(z,xi,zs,z_min,z_max)
        import r_kind
        real(r_kind), intent(out) :: z
        real(r_kind), intent(in ) :: xi
        real(r_kind), intent(in ) :: zs
        real(r_kind), intent(in ) :: z_min
        real(r_kind), intent(in ) :: z_max
      end subroutine vertical_coordinate_interface
      
      subroutine vertical_coordinate_deriv_interface(dzdeta,dzdx,dzdy,xi,zs,dzsdx,dzsdy,dxideta,detadxi,z_min,z_max)
        import r_kind
        real(r_kind), intent(out) :: dzdeta
        real(r_kind), intent(out) :: dzdx
        real(r_kind), intent(out) :: dzdy
        real(r_kind), intent(in ) :: xi
        real(r_kind), intent(in ) :: zs
        real(r_kind), intent(in ) :: dzsdx
        real(r_kind), intent(in ) :: dzsdy
        real(r_kind), intent(in ) :: dxideta
        real(r_kind), intent(in ) :: detadxi
        real(r_kind), intent(in ) :: z_min
        real(r_kind), intent(in ) :: z_max
      end subroutine vertical_coordinate_deriv_interface
    end interface
    
    procedure(vertical_distribution_interface    ), pointer :: vertical_distribution     => null()
    procedure(vertical_coordinate_interface      ), pointer :: vertical_coordinate       => null()
    procedure(vertical_coordinate_deriv_interface), pointer :: vertical_coordinate_deriv => null()
    
    contains
    
    subroutine vertical_init(static)
      type(static_type), intent(inout) :: static
      
      real(r_kind) :: dxideta, detadxi
      
      integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge
      integer idom, ims, ime, jms, jme, kms, kme
      integer i, j, k, p
      
      call log_notice('Initializing vertical distribution and coordinate.', pid=proc%id)

      select case(v_dist_type)
      case(1)
        vertical_distribution => vertical_distribution_equal
      case(2)
        vertical_distribution => vertical_distribution_tan
      case(3)
        vertical_distribution => vertical_distribution_tanh
      end select
      
      select case(v_coord_type)
      case(1)
        vertical_coordinate => vertical_coordinate_gal_chen
      case(2)               
        vertical_coordinate => vertical_coordinate_schar
      end select
      
      associate( mesh => static%mesh )
        do idom = 1, mesh%ndom
          if (mesh%domain(idom)%initialized) then
            associate (ghost   => mesh%domain(idom)%ghost, &
                       z       => mesh%domain(idom)%z,     &
                       eta     => mesh%domain(idom)%eta,   &
                       deta    => mesh%domain(idom)%deta,  &
                       xi      => mesh%domain(idom)%xi,    &
                       eta_min => mesh%domain(idom)%etamin,& 
                       eta_max => mesh%domain(idom)%etamax,& 
                       z_min   => mesh%domain(idom)%zmin,  &
                       z_max   => mesh%domain(idom)%zmax,  &
                       zs      => static%domain(idom)%zs,  &
                       zsC     => static%domain(idom)%zsC, &
                       zsL     => static%domain(idom)%zsL, &
                       zsR     => static%domain(idom)%zsR, &
                       zsB     => static%domain(idom)%zsB, &
                       zsT     => static%domain(idom)%zsT, &
                       zsU     => static%domain(idom)%zsU, &
                       zsD     => static%domain(idom)%zsD, &
                       zsQ     => static%domain(idom)%zsQ  )
            call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme, &
                                 npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
            eta_min = 0
            eta_max = max_height
            deta    = ( eta_max - eta_min ) / nz
            eta     = eta * deta
            z_min   = 0
            z_max   = max_height
            !$OMP PARALLEL DO PRIVATE(i,j,p,dxideta,detadxi) COLLAPSE(3)
            do k = kms, kme
              do j = jms, jme
                do i = ims, ime
                  !z(:,i,j,k) = 0
                  
                  ! Center
                  do p = 1,pc
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsC(i,j),z_min,z_max)
                  end do
                  z(pvs:pve,i,j,k) = z(pc,i,j,k)
                  
                  ! Bottom
                  do p = pes(3), pee(3)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsB(p,i,j),z_min,z_max)
                  end do
                  ! Right
                  do p = pes(2), pee(2)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsR(p,i,j),z_min,z_max)
                  end do
                  ! Top
                  do p = pes(4), pee(4)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsT(p,i,j),z_min,z_max)
                  end do
                  ! Left
                  do p = pes(1), pee(1)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsL(p,i,j),z_min,z_max)
                  end do
                  ! Down
                  do p = pes(5), pee(5)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsD(p,i,j),z_min,z_max)
                  end do
                  ! Up
                  do p = pes(6), pee(6)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsU(p,i,j),z_min,z_max)
                  end do
                  
                  ! Quadrature points
                  do p = pqs, pqe
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(z(p,i,j,k),xi(p,k),zsQ(p,i,j),z_min,z_max)
                  end do
                  
                  ! Ghost points
                  do p = 1, mesh%domain(idom)%ghost(i,j)%ngp
                    ghost(i,j)%eta(p,k) = ghost(i,j)%eta(p,k) * deta
                    call vertical_distribution(ghost(i,j)%xi(p,k),dxideta,detadxi,ghost(i,j)%eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate(ghost(i,j)%z(p,k),ghost(i,j)%xi(p,k),ghost(i,j)%zs(p),z_min,z_max)
                  end do
                end do
              end do
            end do
            !$OMP END PARALLEL DO
            end associate
          endif
        end do
      end associate
    end subroutine vertical_init
    
    subroutine vertical_deriv_init(static)
      type(static_type), intent(inout) :: static
      
      real(r_kind) :: dxideta, detadxi, Ra, Rb, Rr
      
      integer npt, pc, pvs, pve, pes(6), pee(6), pqs, pqe, pgs, pge
      integer idom, ims, ime, jms, jme, kms, kme, igs, ige, jgs, jge
      integer i, j, k, p

      call log_notice('Initializing vertical derivative.', pid=proc%id)
      
      select case(v_coord_type)
      case(1)
        vertical_coordinate_deriv => vertical_coordinate_deriv_gal_chen
      case(2)
        vertical_coordinate_deriv => vertical_coordinate_deriv_schar
      end select
      
      associate( mesh => static%mesh )
        do idom = 1, mesh%ndom
          if (mesh%domain(idom)%initialized) then
            associate (ghost    => mesh%domain(idom)%ghost,    &
                       eta      => mesh%domain(idom)%eta,      &
                       xi       => mesh%domain(idom)%xi,       &
                       eta_min  => mesh%domain(idom)%etamin,   &
                       eta_max  => mesh%domain(idom)%etamax,   &
                       z_min    => mesh%domain(idom)%zmin,     &
                       z_max    => mesh%domain(idom)%zmax,     &
                       zsC      => static%domain(idom)%zsC,    &  
                       zsL      => static%domain(idom)%zsL,    &  
                       zsR      => static%domain(idom)%zsR,    &  
                       zsB      => static%domain(idom)%zsB,    &  
                       zsT      => static%domain(idom)%zsT,    &  
                       zsU      => static%domain(idom)%zsU,    &  
                       zsD      => static%domain(idom)%zsD,    &   
                       zsQ      => static%domain(idom)%zsQ,    &    
                       dzsdx    => static%domain(idom)%dzsdx,  &
                       dzsdxC   => static%domain(idom)%dzsdxC, &
                       dzsdxL   => static%domain(idom)%dzsdxL, &
                       dzsdxR   => static%domain(idom)%dzsdxR, &
                       dzsdxB   => static%domain(idom)%dzsdxB, &
                       dzsdxT   => static%domain(idom)%dzsdxT, &
                       dzsdxU   => static%domain(idom)%dzsdxU, &
                       dzsdxD   => static%domain(idom)%dzsdxD, &
                       dzsdxQ   => static%domain(idom)%dzsdxQ, &
                       dzsdy    => static%domain(idom)%dzsdy,  &
                       dzsdyC   => static%domain(idom)%dzsdyC, &
                       dzsdyL   => static%domain(idom)%dzsdyL, &
                       dzsdyR   => static%domain(idom)%dzsdyR, &
                       dzsdyB   => static%domain(idom)%dzsdyB, &
                       dzsdyT   => static%domain(idom)%dzsdyT, &
                       dzsdyU   => static%domain(idom)%dzsdyU, &
                       dzsdyD   => static%domain(idom)%dzsdyD, &
                       dzsdyQ   => static%domain(idom)%dzsdyQ  )
            call mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme, igs=igs, ige=ige, jgs=jgs, jge=jge,&
                                 npt=npt, pc=pc, pvs=pvs, pve=pve, pes=pes, pee=pee, pqs=pqs, pqe=pqe, pgs=pgs, pge=pge)
            !$OMP PARALLEL
            !$OMP DO PRIVATE(p,i,j,dxideta,detadxi,Ra,Rb,Rr) COLLAPSE(3)
            do k = kms, kme
              do j = jms, jme
                do i = ims, ime
                  ! Center
                  do p = 1,pc
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsC(i,j),dzsdxC(i,j),dzsdyC(i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  
                  ! Bottom
                  do p = pes(3), pee(3)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsB(p,i,j),dzsdxB(p,i,j),dzsdyB(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  ! Right
                  do p = pes(2), pee(2)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsR(p,i,j),dzsdxR(p,i,j),dzsdyR(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  ! Top
                  do p = pes(4), pee(4)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsT(p,i,j),dzsdxT(p,i,j),dzsdyT(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  ! Left
                  do p = pes(1), pee(1)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsL(p,i,j),dzsdxL(p,i,j),dzsdyL(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  ! Down
                  do p = pes(5), pee(5)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsD(p,i,j),dzsdxD(p,i,j),dzsdyD(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  ! Up
                  do p = pes(6), pee(6)
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsU(p,i,j),dzsdxU(p,i,j),dzsdyU(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                  
                  ! Quadrature points
                  do p = pqs, pqe
                    call vertical_distribution(xi(p,k),dxideta,detadxi,eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra,Rb,Rr,xi(p,k),zsQ(p,i,j),dzsdxQ(p,i,j),dzsdyQ(p,i,j),dxideta,detadxi,z_min,z_max)
                    call mesh%domain(idom)%set_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                end do
              end do
            end do
            !$OMP END DO!! NOWAIT

            ! Ghost points
            !$OMP DO PRIVATE(p,i,j,dxideta,detadxi,Ra,Rb,Rr) COLLAPSE(3)
            do k = kms, kme
              do j = jgs, jge
                do i = igs, ige
                  do p = 1, ghost(i,j)%ngp
                    call vertical_distribution(ghost(i,j)%xi(p,k),dxideta,detadxi,ghost(i,j)%eta(p,k),eta_min,eta_max,z_min,z_max,m_coef)
                    call vertical_coordinate_deriv(Ra, Rb, Rr, ghost(i,j)%xi(p,k), ghost(i,j)%zs(p), ghost(i,j)%dzsdx(p), ghost(i,j)%dzsdy(p),&
                                                   dxideta, detadxi, z_min, z_max)
                    call mesh%domain(idom)%set_ghost_ver_metrics(Ra,Rb,Rr,p,i,j,k)
                  end do
                enddo
              enddo
            enddo
            !$OMP END DO
            !$OMP END PARALLEL
            end associate
          endif
        end do
      end associate
    end subroutine vertical_deriv_init
    
    subroutine vertical_distribution_equal(xi,dxideta,detadxi,eta,eta_min,eta_max,z_min,z_max,m_coef)
      real(r_kind), intent(out) :: xi
      real(r_kind), intent(out) :: dxideta
      real(r_kind), intent(out) :: detadxi
      real(r_kind), intent(in ) :: eta
      real(r_kind), intent(in ) :: eta_min,eta_max
      real(r_kind), intent(in ) :: z_min,z_max
      real(r_kind), intent(in ) :: m_coef
      
      real(r_kind) dxi, deta
      
      dxi  = ( z_max - z_min ) / nz
      deta = ( eta_max - eta_min ) / nz
      
      xi = eta * dxi / deta
      
      dxideta = dxi / deta
      detadxi = 1. / dxi
    
    end subroutine vertical_distribution_equal
    
    subroutine vertical_distribution_tan(xi,dxideta,detadxi,eta,eta_min,eta_max,z_min,z_max,m_coef)
      real(r_kind), intent(out) :: xi
      real(r_kind), intent(out) :: dxideta
      real(r_kind), intent(out) :: detadxi
      real(r_kind), intent(in ) :: eta
      real(r_kind), intent(in ) :: eta_min,eta_max
      real(r_kind), intent(in ) :: z_min,z_max
      real(r_kind), intent(in ) :: m_coef
      
      real(r_kind) :: eta_x, x, x_min, x_max
      
      eta_x   = eta / eta_max
      x       = 2. * pi * m_coef * ( eta_x - 0.5 )
      x_min   = -pi * m_coef
      x_max   = x
      xi      = z_max / ( pi**2 * m_coef ) * ( x_max * atan(x_max) - x_min * atan(x_min) + log( sqrt( (1.+x_min**2) / (1.+x_max**2) ) ) ) + z_max * eta_x
      dxideta = 2. * z_max / pi * ( atan(x) + pi / 2. ) / eta_max
      detadxi = 1. / dxideta
    
    end subroutine vertical_distribution_tan
    
    subroutine vertical_distribution_tanh(xi,dxideta,detadxi,eta,eta_min,eta_max,z_min,z_max,m_coef)
      real(r_kind), intent(out) :: xi
      real(r_kind), intent(out) :: dxideta
      real(r_kind), intent(out) :: detadxi
      real(r_kind), intent(in ) :: eta
      real(r_kind), intent(in ) :: eta_min,eta_max
      real(r_kind), intent(in ) :: z_min,z_max
      real(r_kind), intent(in ) :: m_coef
      
      real(r_kind) :: eta_x, x, x_min, x_max, coef, dxh
      
      x_min   = - 0.5 * m_coef
      x_max   = 0.5 * m_coef
      dxh     = x_max - x_min
      coef    = 0.5 / x_max * z_max
      eta_x   = m_coef / eta_max
      x       = ( eta / eta_max - 0.5 ) * m_coef
      xi      = coef * ( ( log( cosh(x) ) - log( cosh(x_min) ) ) + x + dxh / 2. )
      dxideta = coef * ( tanh(x) + 1. ) * eta_x
      detadxi = 1. / dxideta
    
    end subroutine vertical_distribution_tanh
      
    ! Gal-Chen, 1975
    subroutine vertical_coordinate_gal_chen(z,xi,zs,z_min,z_max)
      real(r_kind), intent(out) :: z
      real(r_kind), intent(in ) :: xi
      real(r_kind), intent(in ) :: zs
      real(r_kind), intent(in ) :: z_min
      real(r_kind), intent(in ) :: z_max
      
      z = ( z_max - zs ) / z_max * xi + zs
    
    end subroutine vertical_coordinate_gal_chen
      
    ! Schar, 2002
    subroutine vertical_coordinate_schar(z,xi,zs,z_min,z_max)
      real(r_kind), intent(out) :: z
      real(r_kind), intent(in ) :: xi
      real(r_kind), intent(in ) :: zs
      real(r_kind), intent(in ) :: z_min
      real(r_kind), intent(in ) :: z_max
      
      real(r_kind) :: H, s
      
      H = z_max - z_min
      s = H / exp(1.)
      z = xi + zs * sinh( ( H - xi ) / s ) / sinh( H / s )
    end subroutine vertical_coordinate_schar
      
    ! Gal-Chen, 1975
    subroutine vertical_coordinate_deriv_gal_chen(dzdx,dzdy,dzdeta,xi,zs,dzsdx,dzsdy,dxideta,detadxi,z_min,z_max)
      real(r_kind), intent(out) :: dzdx
      real(r_kind), intent(out) :: dzdy
      real(r_kind), intent(out) :: dzdeta
      real(r_kind), intent(in ) :: xi
      real(r_kind), intent(in ) :: zs
      real(r_kind), intent(in ) :: dzsdx
      real(r_kind), intent(in ) :: dzsdy
      real(r_kind), intent(in ) :: dxideta
      real(r_kind), intent(in ) :: detadxi
      real(r_kind), intent(in ) :: z_min
      real(r_kind), intent(in ) :: z_max
      
      real(r_kind) :: dzdxi, dxidx, dxidy
    
      dzdxi  = ( z_max - zs ) / z_max
      dzdx   = ( z_max - xi ) / z_max * dzsdx
      dzdy   = ( z_max - xi ) / z_max * dzsdy
      dzdeta = dzdxi * dxideta
      
    end subroutine vertical_coordinate_deriv_gal_chen
      
    ! Schar, 2002
    subroutine vertical_coordinate_deriv_schar(dzdx,dzdy,dzdeta,xi,zs,dzsdx,dzsdy,dxideta,detadxi,z_min,z_max)
      real(r_kind), intent(out) :: dzdx
      real(r_kind), intent(out) :: dzdy
      real(r_kind), intent(out) :: dzdeta
      real(r_kind), intent(in ) :: xi
      real(r_kind), intent(in ) :: zs
      real(r_kind), intent(in ) :: dzsdx
      real(r_kind), intent(in ) :: dzsdy
      real(r_kind), intent(in ) :: dxideta
      real(r_kind), intent(in ) :: detadxi
      real(r_kind), intent(in ) :: z_min
      real(r_kind), intent(in ) :: z_max
      
      real(r_kind) :: H, s
      real(r_kind) :: dzdxi, dxidx, dxidy
      
      H = z_max - z_min
      s = H / exp(1._r_kind)
      
      dzdxi  = 1. - zs * cosh( ( H - xi ) / s ) / ( s * sinh( H / s ) )
      dzdx   = dzsdx * sinh( ( H - xi ) / s ) / sinh( H / s )
      dzdy   = dzsdy * sinh( ( H - xi ) / s ) / sinh( H / s )
      dzdeta = dzdxi * dxideta
      
    end subroutine vertical_coordinate_deriv_schar
    
  end module vertical_coord_mod